import torch
from torch import nn

from utils.iou import calculate_iou


# def get_loss(pred, labels):
#     lambda_coord = 5
#

class Loss_yolov1(nn.Module):
    def __init__(self):
        super(Loss_yolov1, self).__init__()

    def forward(self, pred, labels):
        """
        :param pred: (batchsize,30,7,7)的网络输出数据
        :param labels: (batchsize,30,7,7)的样本标签数据
        :return: 当前批次样本的平均损失
        """
        num_gridx, num_gridy = labels.size()[-2:]  # 划分网格数量
        # num_b = 2  # 每个网格的bbox数量
        # num_cls = 20  # 类别数量
        # noobj_confi_loss = 0.  # 不含目标的网格损失(只有置信度损失)
        # coor_loss = 0.  # 含有目标的bbox的坐标损失
        # obj_confi_loss = 0.  # 含有目标的bbox的置信度损失
        # class_loss = 0.  # 含有目标的网格的类别损失
        n_batch = labels.size()[0]  # batchsize的大小

        center_pos_loss = 0.  # 中心点坐标损失
        width_height_loss = 0.  # 物体宽高损失
        conf_loss = 0.
        p_class_loss = 0.

        # 可以考虑用矩阵运算进行优化，提高速度，为了准确起见，这里还是用循环
        for batch_index in range(n_batch):  # batchsize循环
            for x_index in range(num_gridx):  # y方向网格循环
                for y_index in range(num_gridy):  # x方向网格循环
                    if labels[batch_index, 4, x_index, y_index] == 1:  # 如果包含物体
                        # 将数据(px,py,w,h)转换为(x1,y1,x2,y2)
                        # 先将px,py转换为cx,cy，即相对网格的位置转换为标准化后实际的bbox中心位置cx,xy
                        # 然后再利用(cx-w/2,cy-h/2,cx+w/2,cy+h/2)转换为xyxy形式，用于计算iou

                        # 预测框中心点在区域中的相对位置 => 预测框左上角在整个图像中的相对位置
                        bbox1_pred_xyxy = ((pred[batch_index, 0, x_index, y_index] + y_index) / num_gridx
                                           - pred[batch_index, 2, x_index, y_index] / 2,
                                           (pred[batch_index, 1, x_index, y_index] + x_index) / num_gridy
                                           - pred[batch_index, 3, x_index, y_index] / 2,
                                           (pred[batch_index, 0, x_index, y_index] + y_index) / num_gridx
                                           + pred[batch_index, 2, x_index, y_index] / 2,
                                           (pred[batch_index, 1, x_index, y_index] + x_index) / num_gridy
                                           + pred[batch_index, 3, x_index, y_index] / 2)
                        bbox2_pred_xyxy = ((pred[batch_index, 5, x_index, y_index] + y_index) / num_gridx
                                           - pred[batch_index, 7, x_index, y_index] / 2,
                                           (pred[batch_index, 6, x_index, y_index] + x_index) / num_gridy
                                           - pred[batch_index, 8, x_index, y_index] / 2,
                                           (pred[batch_index, 5, x_index, y_index] + y_index) / num_gridx
                                           + pred[batch_index, 7, x_index, y_index] / 2,
                                           (pred[batch_index, 6, x_index, y_index] + x_index) / num_gridy
                                           + pred[batch_index, 8, x_index, y_index] / 2)
                        bbox_gt_xyxy = ((labels[batch_index, 0, x_index, y_index] + y_index) / num_gridx
                                        - labels[batch_index, 2, x_index, y_index] / 2,
                                        (labels[batch_index, 1, x_index, y_index] + x_index) / num_gridy
                                        - labels[batch_index, 3, x_index, y_index] / 2,
                                        (labels[batch_index, 0, x_index, y_index] + y_index) / num_gridx
                                        + labels[batch_index, 2, x_index, y_index] / 2,
                                        (labels[batch_index, 1, x_index, y_index] + x_index) / num_gridy
                                        + labels[batch_index, 3, x_index, y_index] / 2)
                        iou1 = calculate_iou(bbox1_pred_xyxy, bbox_gt_xyxy)
                        iou2 = calculate_iou(bbox2_pred_xyxy, bbox_gt_xyxy)
                        # 选择iou大的bbox作为负责物体
                        if iou1 >= iou2:
                            # 中心位置损失
                            center_pos_loss += 5 * (torch.sum((pred[batch_index, 0:2, x_index, y_index]
                                                               - labels[batch_index, 0:2, x_index, y_index]) ** 2))
                            width_height_loss += 5 * (
                                torch.sum(
                                    (labels[batch_index, 2, x_index, y_index].sqrt()
                                     - pred[batch_index, 2, x_index, y_index].sqrt()) ** 2  # 宽度损失
                                    + (labels[batch_index, 3, x_index, y_index].sqrt()
                                       - pred[batch_index, 3, x_index, y_index].sqrt()) ** 2  # 高度损失
                                )
                            )
                            conf_loss += (torch.sum((labels[batch_index, 4, x_index, y_index]
                                                     - pred[batch_index, 4, x_index, y_index]) ** 2)  # 正类置信损失
                                          + 0.5 * torch.sum((labels[batch_index, 9, x_index, y_index]
                                                             - pred[batch_index, 9, x_index, y_index]) ** 2)  # 负类置信损失
                                          )

                            p_class_loss += torch.sum(
                                (pred[batch_index, 10:, x_index, y_index]
                                 - labels[batch_index, 10:, x_index, y_index]) ** 2  # 分类损失
                            )
                        else:
                            # 中心位置损失
                            center_pos_loss += 5 * (torch.sum((pred[batch_index, 5:7, x_index, y_index]
                                                               - labels[batch_index, 5:7, x_index, y_index]) ** 2))
                            width_height_loss += 5 * (
                                torch.sum(
                                    (labels[batch_index, 7, x_index, y_index].sqrt()
                                     - pred[batch_index, 7, x_index, y_index].sqrt()) ** 2  # 宽度损失
                                    + (labels[batch_index, 8, x_index, y_index].sqrt()
                                       - pred[batch_index, 8, x_index, y_index].sqrt()) ** 2  # 高度损失
                                )
                            )
                            conf_loss += (torch.sum((labels[batch_index, 9, x_index, y_index]
                                                     - pred[batch_index, 9, x_index, y_index]) ** 2)  # 正类置信损失
                                          + 0.5 * torch.sum((labels[batch_index, 4, x_index, y_index]
                                                             - pred[batch_index, 4, x_index, y_index]) ** 2)  # 负类置信损失
                                          )

                            p_class_loss += torch.sum(
                                (pred[batch_index, 10:, x_index, y_index]
                                 - labels[batch_index, 10:, x_index, y_index]) ** 2  # 分类损失
                            )

                    else:  # 如果不包含物体
                        conf_loss += 0.5 * torch.sum((pred[batch_index, [4, 9], x_index, y_index]) ** 2)  # 负类置信损失
                        # noobj_confi_loss = noobj_confi_loss + 0.5 * torch.sum(
                        #     pred[batch_index, [4, 9], x_index, y_index] ** 2)

        all_loss = center_pos_loss + width_height_loss + conf_loss + p_class_loss
        # loss = coor_loss + obj_confi_loss + noobj_confi_loss + class_loss
        # 此处可以写代码验证一下loss的大致计算是否正确，这个要验证起来比较麻烦，比较简洁的办法是，将输入的pred置为全1矩阵，再进行误差检查，会直观很多。
        return all_loss / n_batch

        # return loss / n_batch
